package com.lyx.rpc.registry.zookeeper;

import com.alibaba.fastjson.serializer.JSONSerializableSerializer;
import com.lyx.rpc.common.helper.RpcServiceHelper;
import com.lyx.rpc.loadbalancer.api.ServiceLoadBalancer;
import com.lyx.rpc.loadbalancer.random.RandomServiceLoadBalancer;
import com.lyx.rpc.protocol.meta.ServiceMeta;
import com.lyx.rpc.registry.api.RegistryService;
import com.lyx.rpc.registry.api.config.RegistryConfig;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Random;

/**
 * @Author by lyx
 * 基于Zookeeper的注册服务
 */
public class ZookeeperRegistryService implements RegistryService {
    /**
     * 初始化CuratorFramework客户端时，进行连接重试的间隔时间。
     */
    public static final int BASE_SLEEP_TIME_MS = 1000;
    /**
     * 初始化CuratorFramework客户端时，进行连接重试的最大重试次数。
     */
    public static final int MAX_RETRIES = 3;
    /**
     * 服务注册到Zookeeper的根路径，我们写的RPC框架的根目录为/lyx_rpc。
     */
    public static final String ZK_BASE_PATH = "/lyx_rpc";
    private static final Logger log = LoggerFactory.getLogger(ZookeeperRegistryService.class);
    /**
     * 服务注册与发现的ServiceDiscovery类实例。
     */
    private ServiceDiscovery<ServiceMeta> serviceDiscovery;
    private ServiceLoadBalancer<ServiceInstance<ServiceMeta>> serviceLoadBalancer;

    @Override
    public void register(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> instance = ServiceInstance.<ServiceMeta>builder()
            .name(RpcServiceHelper.buildServiceKey(serviceMeta.getServiceName(), serviceMeta.getServiceVersion(), serviceMeta.getServiceGroup()))
            .address(serviceMeta.getServiceAddr())
            .port(serviceMeta.getPort())
            .payload(serviceMeta)
            .build();
        this.serviceDiscovery.registerService(instance);
    }

    @Override
    public void unRegister(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> instance = ServiceInstance.<ServiceMeta>builder()
            .name(RpcServiceHelper.buildServiceKey(serviceMeta.getServiceName(), serviceMeta.getServiceVersion(), serviceMeta.getServiceGroup()))
            .address(serviceMeta.getServiceAddr())
            .port(serviceMeta.getPort())
            .payload(serviceMeta)
            .build();
        this.serviceDiscovery.unregisterService(instance);
    }

    @Override
    public ServiceMeta discovery(String serviceName, int invokerHashCode) throws Exception {
        Collection<ServiceInstance<ServiceMeta>> serviceInstances = this.serviceDiscovery.queryForInstances(serviceName);
        ServiceInstance<ServiceMeta> instance = this.serviceLoadBalancer.select((List<ServiceInstance<ServiceMeta>>) serviceInstances, invokerHashCode);
        if (instance != null) {
            log.info(String.valueOf(instance));
            return instance.getPayload();
        }
        return null;
    }

    @Override
    public void destroy() throws IOException {
        this.serviceDiscovery.close();
    }

    @Override
    public void init(RegistryConfig registryConfig) throws Exception {
        CuratorFramework client =
            CuratorFrameworkFactory.newClient(registryConfig.getRegistryAddr(), new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();
        JsonInstanceSerializer<ServiceMeta> serializer = new JsonInstanceSerializer<>(ServiceMeta.class);
        this.serviceDiscovery =
            ServiceDiscoveryBuilder.builder(ServiceMeta.class)
                .serializer(serializer)
                .client(client)
                .basePath(ZK_BASE_PATH).build();
        this.serviceDiscovery.start();
        this.serviceLoadBalancer = new RandomServiceLoadBalancer<ServiceInstance<ServiceMeta>>();
    }
}
